#!/usr/bin/python
import sys, os, re, time, math
################################################################################
################################################################################
class Params:
   def __init__( self ):
      pass
   

class MtxInfo:
   m   = 0
   n   = 0
   nnz = 0
   bytes = 0
   bytesExpanded = 0
   symmetric = False
   def ParseInfoLine( self, infoline ):
      toks = infoline.split()
      assert toks[0] == 'MatrixMarket'
      assert toks[1] == 'matrix'
      assert toks[2] == 'coordinate'
      assert toks[3] == 'real'
      assert toks[4] == 'general' or toks[4] == 'symmetric'
      if toks[4] == 'symmetric':
         self.symmetric = True
      else:
         self.symmetric = False
   
   def CalculateMtxBytes( self ):
      sizeDoubles = 8
      sizeIndices = 4
      sizePointer = 4
      numColIdxs = self.nnz
      numRowPtrs = self.m + 1
      bytesDoubles = sizeDoubles * self.nnz
      bytesColIdxs = sizeIndices * numColIdxs
      bytesRowPtrs = sizePointer * numRowPtrs
      self.bytes = bytesDoubles + bytesColIdxs + bytesRowPtrs
      self.bytes = float( self.bytes ) / (1024.0 * 1024.0)
      if self.symmetric:
         assert self.m == self.n
         totnnz = self.nnz + ( self.nnz - self.m )
         numColIdxs = totnnz
         numRowPtrs = self.m + 1
         bytesDoubles = sizeDoubles * totnnz
         bytesColIdxs = sizeIndices * numColIdxs
         bytesRowPtrs = sizePointer * numRowPtrs
         self.bytesExpanded = bytesDoubles + bytesColIdxs + bytesRowPtrs
         self.bytesExpanded = float( self.bytesExpanded ) / (1024.0 * 1024.0)
   
   def __init__( self, fn ):
      dimsLine = MtxFileDimsLine( fn )
      infoLine = MtxFileInfoLine( fn )
      toks = dimsLine.split()
      self.m   = int( toks[0] )
      self.n   = int( toks[1] )
      self.nnz = int( toks[2] )   
      self.ParseInfoLine( infoLine )
      self.CalculateMtxBytes()
   

class DagInfo:
   m   = 0
   n   = 0
   deg = 0
   nnz = 0
   numNodes = 0
   bytes = 0
   def CalculateDagBytes( self ):
      bytesPerNode = 4 * self.deg
      self.bytes = bytesPerNode * self.numNodes
      self.bytes = float( self.bytes ) / (1024.0 * 1024.0)
   
   def __init__( self, fn ):
      toplines = HeadNLines( fn, 100 )
      self.m   = int( toplines[0] )
      self.n   = int( toplines[1] )
      self.nnz = int( toplines[2] )
      self.numNodes = int( toplines[3] )
      self.deg = 2 * len( toplines[-1].split() )
      self.CalculateDagBytes()
   

################################################################################
################################################################################
def SystemWrapperPipelined( cmd, verbose, trial = False ):
  if verbose:
    PrintTimeStamp()
    print '... os.popen(', cmd, ')'
    sys.stdout.flush()
  if not trial:
    FH = os.popen( cmd )
    output = [ line.rstrip() for line in FH.readlines() ]
    FH.close()
    return output
  return []

################################################################################
################################################################################
def HeadNLines( fn, n ):
   cmd = ' '.join([ 'head', '-%d' % n, fn ])
   return SystemWrapperPipelined( cmd, False, False )

def MtxFileDimsLine( mtxFile ):
   lines = HeadNLines( mtxFile, 100 )
   for line in lines:
      if len( line ) == 0: continue
      if line[0] == '%': continue
      assert len( line.split() ) == 3
      return line
   assert False

def MtxFileInfoLine( mtxFile ):
   infoline = HeadNLines( mtxFile, 1 )[0].rstrip().lstrip( '%' )
   assert len( infoline.split() ) == 5
   return infoline

################################################################################
################################################################################
def ConvertDagFileToMtxFileName( dagfile ):
   ( path, rawf ) = os.path.split( dagfile )
   toks = rawf.split( '.' )
   rawf = '.'.join([ toks[0], 'mtx' ])
   return os.path.join( path, rawf )

################################################################################
################################################################################
def PrintComparison( mtxInfo, dagInfo ):
   # print mtxInfo.m
   # print mtxInfo.n
   # print mtxInfo.nnz
   # print mtxInfo.symmetric
   # print mtxInfo.bytes / ( 1024*1024)
   # print mtxInfo.bytesExpanded / ( 1024*1024)
   # print '-'
   # print dagInfo.m
   # print dagInfo.n
   # print dagInfo.nnz
   # print dagInfo.bytes  / ( 1024*1024)
   cr = float( mtxInfo.bytes ) / float( dagInfo.bytes )
   cr = '%.1f ' % cr
   fr = '%.1f'  % mtxInfo.bytes
   to = '%.1f'  % dagInfo.bytes
   print '%5sM =>%5sM, %5sx' % ( fr, to, cr )
   if mtxInfo.symmetric:
      cr = float( mtxInfo.bytesExpanded ) / float( dagInfo.bytes )
      cr = '%.1f'  % cr
      fr = '%.1f'  % mtxInfo.bytesExpanded
      to = '%.1f'  % dagInfo.bytes
      print '%5sM =>%5sM, %5sx' % ( fr, to, cr )

################################################################################
################################################################################
def CheckArgs( argv ):
   params = Params()
   argc = len( argv )
   if argc < 2:
      print 'Usage:', argv[0], '<dag_file>'
      sys.exit( -1 )
   
   params.dagfile = argv[1]
   if not os.path.isfile( params.dagfile ):
      print params.dagfile, 'is not a file.'
      print 'Must at least supply a valid filename'
      sys.exit( -1 )
      
   params.mtxfile = ConvertDagFileToMtxFileName( params.dagfile )
   if not os.path.isfile( params.mtxfile ):
      print params.mtxfile, 'is not a file.'
      print 'Must place the original .mtx file in the same place as the dag file'
      sys.exit( -1 )
   
   return params

################################################################################
################################################################################
params = CheckArgs( sys.argv )

mtxInfo = MtxInfo( params.mtxfile )
dagInfo = DagInfo( params.dagfile )

PrintComparison( mtxInfo, dagInfo )
